/*
 * Copyright 2011 TTOT
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package ttot.loaders;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.regex.Pattern;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;

/**
 * Custom classloader for loading class files and resources from jar packages.
 * 
 * @author george.aprozeanu@gmail.com
 * 
 */
public class JarClassLoader extends ClassLoader {

	private HashMap<File, ByteBuffer> contents = new HashMap<File, ByteBuffer>();

	private ArrayList<File> files = new ArrayList<File>();
	private HashMap<String, Class<?>> classes = new HashMap<String, Class<?>>();
	private HashMap<String, ByteBuffer> classData = new HashMap<String, ByteBuffer>();
	private ArrayList<File> directories = new ArrayList<File>();

	public JarClassLoader(ClassLoader classLoader) {
		super(classLoader);
	}

	/**
	 * Reads contents of the given input stream and caches the classes,
	 * resources and folder entries.
	 * 
	 * @param jarContents
	 *            - an input stream that provides the contents on a jar file
	 * @throws IOException
	 *             - cascaded IOException while handling the input stream
	 */
	public void loadJar(InputStream jarContents) throws IOException {
		readStream(jarContents);
	}

	@Override
	protected void finalize() throws Throwable {
		this.files = null;
		this.classes = null;
		this.classData = null;
		this.directories = null;
		super.finalize();
	}

	@Override
	protected Class<?> findClass(String name) throws ClassNotFoundException {
		if (classes.containsKey(name)) {
			return classes.get(name);
		}
		if (classData.containsKey(name)) {
			classes.put(name, defineClass(name, classData.get(name), null));
			return classes.get(name);
		}
		return getParentClassLoader().loadClass(name);
	}

	private ByteBuffer readZipEntry(ZipInputStream zipInputStream) throws IOException {
		ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
		byte[] bytes = new byte[0x4000];
		int readsize;
		while ((readsize = zipInputStream.read(bytes)) != -1) {
			byteArrayOutputStream.write(bytes, 0, readsize);
		}
		byteArrayOutputStream.flush();
		byteArrayOutputStream.close();
		return ByteBuffer.wrap(byteArrayOutputStream.toByteArray());
	}

	private ClassLoader getParentClassLoader() {
		ClassLoader parent = getParent();
		if (parent == null) {
			return getSystemClassLoader();
		}
		return parent;
	}

	@Override
	public Class<?> loadClass(String name) throws ClassNotFoundException {
		return findClass(name);
	}

	private static String classNameOf(File file) {
		String path = file.getPath();
		if (!path.endsWith(".class")) {
			throw new IllegalArgumentException("Invalid class file name");
		}
		String strippedExtension = path.substring(0, path.length() - 6);
		return strippedExtension.replace(File.separatorChar, '.');
	}

	private void registerClass(File file, ByteBuffer data) {
		String className = classNameOf(file);
		try {
			classes.put(className, getParentClassLoader().loadClass(className));
		} catch (ClassNotFoundException e) {
			classData.put(className, data);
		}
	}

	private void readStream(InputStream jarContents) throws IOException {
		ZipInputStream zipInputStream = new ZipInputStream(jarContents);
		ZipEntry zipEntry;
		while ((zipEntry = zipInputStream.getNextEntry()) != null) {
			File entry = new File(zipEntry.getName());
			if (!zipEntry.isDirectory()) {
				ByteBuffer fileContents = readZipEntry(zipInputStream);
				zipInputStream.closeEntry();
				contents.put(entry, fileContents);
				files.add(entry);
				if (isClassFile(entry)) {
					registerClass(entry, fileContents);
				}
			} else {
				directories.add(entry);
			}
		}
		zipInputStream.close();
	}

	private static Pattern classNamePattern = Pattern.compile("^.*[.]class$");

	private static boolean isClassFile(File file) {
		return classNamePattern.matcher(file.getName()).matches();
	}

	/**
	 * @return Class array with all classes from the cached jar stream
	 */
	public Class<?>[] getClasses() {
		Class<?>[] entriesArray = new Class<?>[classes.size()];
		for (String name : classData.keySet()) {
			try {
				loadClass(name);
			} catch (ClassNotFoundException e) {
				throw new RuntimeException(e);
			}
		}
		return classes.values().toArray(entriesArray);
	}

	public File[] getFiles() {
		File[] entriesArray = new File[files.size()];
		return files.toArray(entriesArray);
	}

	public Boolean fileExists(File file) {
		return files.contains(file);
	}

	public File[] getFolders() {
		File[] entriesArray = new File[directories.size()];
		return directories.toArray(entriesArray);
	}

	public ByteBuffer getFileContent(File file) {
		if (contents.containsKey(file)) {
			return contents.get(file);
		} else {
			throw new IllegalArgumentException("Unable to read file content from path " + file.getAbsolutePath());
		}
	}

	public InputStream getResourceAsStream(String path) {
		return new ByteArrayInputStream(getFileContent(new File(path)).array());
	}

}
